#Covered by GPL V2.0
#Coded by Carlos del Ojo Elias (deepbit@gmail.com)


from fuzzEngine import *
import threading
import zlib

class wfDomain:
	def __init__(self):
		self.fuzzs={}
		self.Semaphore_Mutex=threading.BoundedSemaphore(value=1)


	def appendFuzz (self,fuzz):
		self.fuzzs[fuzz.id]=fuzz
		fuzz.Launch()

	def getEncodings (self):
		list=['None']
		list+=ENCODERS.keys()
		return list

	def addPathFuzz(self,ck,headers,proxy,threads,authtype,auth,urls,recursion,extensions,dic):
		ck=ck.replace('\0',"")
		exts=extensions.replace('\0',"").split(",")
		RecursLevel=int(recursion)

		urls=urls.replace('\0',"").split("\n")
		
		if ":" not in proxy:
			proxy=None

		if authtype=="none":
			Auth=[None,None]
		else:
			Auth=[authtype,auth]

		p=payload_file(dic) 
		
		rq=Request()
		if ck:
			rq.addHeader("Cookie",ck)
		
		for i,j in headers.items():
			rq.addHeader(i,j)


		cnt=0
		for i in urls:
			i=i.strip()
			if i:
				cnt+=1
				rq.setUrl(i)
				dic=dictionary()
				dic.setPayload(p)

				self.appendFuzz(FuzzObject(PathFuzzer(rq,dic,maxLevel=RecursLevel,proxy=proxy,threads=threads,Auth=Auth,extensions=exts,saveMemory=False),i))

		return cnt


	def addGeneralFuzz(self,ck,headers,proxy,threads,authtype,auth,url,post,payloads):

		ck=ck.replace('\0',"")
		post=post.replace('\0',"").strip()

		if ":" not in proxy:
			proxy=None

		if authtype=="none":
			Auth=[None,None]
		else:
			Auth=[authtype,auth]

		dicts={}


		for i in payloads:
			id,type,par,enc=i
			if type=="file":
				p=payload_file(par) 
			elif type=="range":
				p=payload_range(par)    
			elif type=="hex-range":
				p=payload_hexrange(par) 
			
			dic=dictionary()
			dic.setPayload(p)
			if enc!="None":
				dic.setEncoder(select_encoding(enc))

			dicts[id]=dic
	

		rq=Request()
		rq.setUrl(url)	
		if ck:
			rq.addHeader("Cookie",ck)
		
		for i,j in headers.items():
			rq.addHeader(i,j)

		if post:
			rq.setPostData(post)

		self.appendFuzz(FuzzObject(Fuzzer(rq,dicts,proxy,threads=threads,FollowLocation=True,Auth=Auth,saveMemory=False),url))
	
	def getAllResults(self,n):
		if n<0:
			return []
		self.Semaphore_Mutex.acquire()
		try:
			r=self.fuzzs[n].getAllResults()
		except:
			r=[]
		self.Semaphore_Mutex.release()
		return r

	def updateResults(self,n):
		if n not in self.fuzzs:
			return {"action":"update","values":[]}
		self.Semaphore_Mutex.acquire()
		r=self.fuzzs[n].updateResults()
		self.Semaphore_Mutex.release()
		return r

	def getResult(self,nfuzz,nres):
		return self.fuzzs[nfuzz].getResult(nres)

	def getinfoFuzz(self):
		list=[]
		for i in self.fuzzs.values():
			list.append(i.description())

		return list

	def cancelFuzz(self,nfuzz):
		self.fuzzs[nfuzz].cancel()
		
	def deleteFuzz(self,nfuzz):
		self.fuzzs[nfuzz].cancel()
		del self.fuzzs[nfuzz]

class FuzzObject:
	ID=0
	def __init__(self,fz,url):
		self.url=url
		self.fz=fz

		self.results=[]
		self.banlist={}
		self.banlistObjects={}

		self.nres=0

		self.id=FuzzObject.ID
		FuzzObject.ID+=1
		
	

	def cancel(self):
		self.fz.stop()
		self.fz.wait()

	def Launch(self):
		self.fz.Launch()

	def getResult(self,nres):
		for i in self.results:
			if i.number==nres:
				return i.content,[i.url,i.headers,i.contenttype]

	def appendContent(self,fuzzResult):
		nfr=FuzzSummary(fuzzResult)


		key=",".join([str(nfr.code),str(nfr.words),str(nfr.lines)])
		if not key in self.banlist:
			self.banlist[key]=0

		if self.banlist[key]<100:
			self.results.append(nfr)
			self.banlist[key]+=1
			if not key in self.banlistObjects:
				self.banlistObjects[key]=[]
			self.banlistObjects[key].append(nfr)

			if self.banlist[key]==100:
				for i in self.banlistObjects[key]:
					self.results.remove(i)
				del self.banlistObjects[key]
				return { "action" : "update", "values": [] }

			return { "action":"append", "values":nfr.summary() }

		return { "action" : "none", "values":[]}


	def getAllResults(self):
		res=self.fz.getAvailableResults()
		for i in res:
			self.nres+=1
			if i.code!=404:
				self.appendContent(i)

		res=[i.summary() for i in self.results]
		return { "action" : "append", "values": res }

	def updateResults(self):
		newresults=[]
		res=self.fz.getAvailableResults()
		update=False

		for i in res:	
			if i.code!=404:
				state=self.appendContent(i)
				if state["action"]=="update":
					update=True
				elif state["action"]=="append":
					newresults.append(state["values"])

		if update:
			return { "action" :"update", "values":[i.summary() for i in self.results] }

		return { "action": "append", "values": newresults }

	def description(self):
		status=self.fz.status()
		return [self.id,self.fz.state,status[0]*100/status[1],"%d/%d"%(status[0],status[1]),status[2],self.url]

class FuzzSummary:
	def __init__(self,fr):
		self.number=fr.number
		self.code=fr.code
		self.lines=fr.lines
		self.words=fr.words
		self.len=fr.len
		self.descrip=fr.descrip.replace("'","&apos;")
		self.descrip=self.descrip.replace("\"","&quot;")

		self.headers,html=fr.getResponse()
		self.__htmlcontent=zlib.compress(html)
		self.url=fr.completeUrl
		self.contenttype=fr.contenttype
	
	def __getattr__ (self,name):
		if name=="content":
			res=zlib.decompress(self.__htmlcontent)
			if not res:
				return "<code><pre>"+self.headers+"</pre></code>"
			return res
		else:
			print name
			raise AttributeError
	
	def summary(self):
		return [self.number,self.code,self.lines,self.words,self.len,self.descrip]

